import {
  Box,
  Collapse,
  Flex,
  FormControl,
  FormHelperText,
  FormLabel,
  Input,
  Radio,
  RadioGroup,
  Stack,
  Textarea,
  useColorModeValue,
  useDisclosure,
} from '@chakra-ui/react';

import { FormEvent, useState } from 'react';
import { MdCheckCircle } from 'react-icons/md';

import { TUploadState } from '.';

import trainService from '@/api/train.service';
import IAIButton from '@/common/components/IAIButton';
import IAISlider from '@/common/components/IAISlider';
import IAISwitch from '@/common/components/IAISwitch';

interface FormProps {
  uploadState: TUploadState;
  afterCreateStudio: () => void;
  ModalType: number | string;
  uploadedImgIds: { [k: string]: string }[];
}
type ImageSize = '512*512' | '512*768' | '768*768';
export default function ParametersForm({
  uploadState,
  afterCreateStudio,
  ModalType,
  uploadedImgIds,
}: FormProps) {
  const [newModalName, setNewModalName] = useState<string>('');
  const [newModalDesc, setNewModalDesc] = useState<string>('');
  const [newModalTrigger, setNewModalTrigger] = useState<string>('');
  const [isLoading, setIsLoading] = useState<boolean>(false);
  const [loraRepeat, setLoraRepeat] = useState<number>(10);
  // const [useAdvanceForm, setUseAdvanceForm] = useState<boolean>(false);
  const [imageSize, setImageSize] = useState<ImageSize>('512*512');
  const { isOpen, onToggle } = useDisclosure();
  const handleSubmitForm = async () => {
    setIsLoading(true);
    const [width, height] = imageSize.split('*');
    try {
      const { code } = await trainService.trainModel({
        name: newModalName,
        modelDesc: newModalDesc,
        triggerWords: newModalTrigger,
        width,
        height,
        imageIds: uploadedImgIds.map((i) => Object.values(i)[0]),
        type: Number(ModalType),
        repeatTime: Number(ModalType) === 2 ? loraRepeat : undefined,
      });
      if (code === 200) {
        setNewModalDesc('');
        setNewModalName('');
        setNewModalTrigger('');
        afterCreateStudio();
      }
    } catch (error) {
      console.log(error);
    }

    setIsLoading(false);
  };

  const formHelperTextColor = useColorModeValue(
    'blackAlpha.600',
    'whiteAlpha.600'
  );

  return (
    <Box
      as="form"
      onSubmit={(e: FormEvent<HTMLDivElement>) => {
        e.preventDefault();
        if (newModalName.trim() && newModalTrigger.trim()) {
          handleSubmitForm();
        }
      }}
      mt={4}
    >
      <Flex gap="8">
        <FormControl isRequired>
          <FormLabel>模型名称</FormLabel>
          <Input
            placeholder="模型名称"
            value={newModalName}
            onChange={(e) => setNewModalName(e.currentTarget.value)}
          />
          {/* <FormHelperText color={formHelperTextColor}>
            建议使用拼音进行命名，并且避免名称重复，用来区分你的模型
          </FormHelperText> */}
        </FormControl>
        <FormControl isRequired>
          <FormLabel>调用触发词</FormLabel>
          <Input
            placeholder="调用触发词"
            value={newModalTrigger}
            onChange={(e) => setNewModalTrigger(e.currentTarget.value)}
          />
          <FormHelperText color={formHelperTextColor}>
            建议使用拼音或英文进行命名，并且避免名称重复，用来区分你的模型
          </FormHelperText>
        </FormControl>
      </Flex>

      <FormControl>
        <FormLabel>模型描述</FormLabel>
        <Textarea
          borderRadius="md"
          value={newModalDesc}
          onChange={(e) => setNewModalDesc(e.currentTarget.value)}
          placeholder="请输入模型的描述"
          size="sm"
        />
        <FormHelperText color={formHelperTextColor}>
          请使用英语对训练素材进行描述
        </FormHelperText>
      </FormControl>

      <Box mt="4">
        <IAISwitch
          width="120px"
          id="advance-switch"
          label="高级参数"
          mt="4px"
          onChange={(e) => {
            // setUseAdvanceForm(e.target.checked);
            onToggle();
          }}
        />
        <Collapse in={isOpen} animateOpacity>
          <Flex gap="8" p="4" flexDirection={{ base: 'column', md: 'row' }}>
            <FormControl w="min-content">
              <FormLabel>图片剪裁尺寸:</FormLabel>
              <RadioGroup
                onChange={(v: ImageSize) => setImageSize(v)}
                value={imageSize}
              >
                <Stack direction="row" spacing="4">
                  <Radio colorScheme="accent" value="512*512">
                    512*512
                  </Radio>
                  <Radio colorScheme="accent" value="512*768">
                    512*768
                  </Radio>
                  <Radio colorScheme="accent" value="768*768">
                    768*768
                  </Radio>
                </Stack>
              </RadioGroup>
              <FormHelperText color={formHelperTextColor}>
                上传的图片将参考原图中心被剪裁为选择的尺寸
              </FormHelperText>
            </FormControl>
            {Number(ModalType) === 2 && (
              <FormControl maxW={{ base: '100%', md: '300px' }}>
                <IAISlider
                  withSliderMarks
                  sliderFormLabelProps={{ mb: 1 }}
                  value={loraRepeat}
                  label="LoRA repeats:"
                  min={10}
                  max={50}
                  step={1}
                  onChange={(value) => {
                    setLoraRepeat(value);
                  }}
                  helperText="训练循环次数，次数越多学习效果越好，但会增加训练时长"
                />
              </FormControl>
            )}
          </Flex>
        </Collapse>
      </Box>

      <Flex my={10} textAlign="center" justify="center">
        <Box maxW="20rem">
          <IAIButton
            rightIcon={<MdCheckCircle />}
            size="lg"
            isLoading={isLoading}
            type="submit"
            isDisabled={uploadState !== 'uploaded'}
          >
            开始训练
          </IAIButton>
        </Box>
      </Flex>
    </Box>
  );
}
